# -*- coding: utf-8 -*-
from __future__ import print_function

import os
import yaml
import time
import subprocess
import datetime
from ddl_platform.ddlib.job import JobStatus

r'''The global GPU resources in the cluster
'''
AVAIL_NODES = ['gpu1', 'gpu2', 'gpu3', 'gpu4']
#AVAIL_NODES = ['gpu1']
#AVAIL_NODES = ['gpu1', 'gpu2', 'gpu3', 'gpu4', 'gpu5', 'gpu6', 'gpu7', 'gpu8', 'gpu10', 'gpu11', 'gpu12', 'gpu13', 'gpu14', 'gpu15', 'gpu16']
#AVAIL_NODES = ['gpu1', 'gpu2', 'gpu3', 'gpu4', 'gpu5', 'gpu6', 'gpu7', 'gpu8', 'gpu10', 'gpu11', 'gpu12', 'gpu13', 'gpu14', 'gpu15', 'gpu16']
#AVAIL_NODES = ['gpu10', 'gpu11', 'gpu12', 'gpu13', 'gpu14', 'gpu15', 'gpu16']
#AVAIL_NODES = ['gpu10']
#AVAIL_NODES = ['scigpu10', 'scigpu11', 'scigpu12', 'scigpu13']

AVAIL_GPUS = {}
for node in AVAIL_NODES:
    AVAIL_GPUS[node] = [0,1,2,3]


class Cluster:
    r'''The cluster receives a file that records the hardware configuration

    Args:
        config (str): The cluster configuration file

    Attributes:
        
    '''

    def __init__(self, config=None): 

        self._nodes = []
        self._node_map = {}
        self._gpu_map = {}
        self._nnodes = len(AVAIL_NODES)
        self._ngpus = 0

        for node_id, node_name in enumerate(AVAIL_NODES):
            gpus = []
            for i in AVAIL_GPUS[node_name]:
                gpu = GPU(node_id, node_name, '%s-%s'%(str(node_id), str(i)), i)
                gpus.append(gpu)
                #gpu_key = '%d-%d'%(node_id, gpu.get_gpu_id())
                self._gpu_map[gpu.get_gpu_id()] = gpu
            #gpus = [GPU(node_id, node_name, i) for i in AVAIL_GPUS[node_name]]
            node = Node(node_id, node_name, gpus)
            self._node_map[node_id] = node
            self._nodes.append(node)
            self._ngpus += len(gpus)

        if config != None:
            pass

    def update_resources(self):
        '''TODO:
           update the cluster resources by profiling
        '''
        pass

    def update_cluster_resources(self, job_list):
        # print('update_cluster_resources..........')
        matched_gpu_job = None
        for job in job_list:
            #print('Current job: ', job)
            node_ids = list(set(job.get_node_ids()))
            occupied_gpus = []
            for node_id in node_ids:
                node = self._node_map[node_id]
                for gpu in node.get_gpus():
                    gpu_job = gpu.get_job(job._job_id)
                    if gpu_job:
                        new_mem = job.get_occupied_mem()
                        gpu_job.update_occupied_mem(new_mem)
                        gpu_job._accumulate_running_time = job._accumulate_running_time
                        gpu_job._bs_iteration_times = job._bs_iteration_times

                        gpu.update_allocated_mem(job._job_id, new_mem)
                        occupied_gpus.append(gpu)
            #print('\tOccupied GPUs:', occupied_gpus)

    def print_resources(self):
        print()
        print(datetime.datetime.now())
        for node in self.get_nodes():
            print(node)
            for gpu in node.get_gpus():
                print(gpu)
                #gpu.print_mem()

    def get_resources_info(self):
        logline = "\n"
        # logline += str(datetime.datetime.now())
        for node in self.get_nodes():
            logline += '%s\n' % str(node)
            for gpu in node.get_gpus():
                logline += '%s\n' % str(gpu)
                #gpu.print_mem()
        return logline

    def get_nodes(self):
        return self._nodes

    def get_gpus(self):
        gpu_list = []
        for node in self._nodes:
            gpu_list.extend(node._gpus)

        return gpu_list

    def get_gpu(self, gpu_id):
        return self._gpu_map.get(gpu_id, None)

    def check_comp(self, node_ids, gpu_ids):
        isFree = True
        for nid in node_ids:
            for gid in gpu_ids:
                isFree = self._nodes[nid].get_gpus()[gid].get_status()
            if not isFree:
                return False
        return True

    def check_comm(self, node_ids):
        isFree = True
        for nid in node_ids:
            isFree = self._nodes[nid].get_status()
            if not isFree:
                return False
        return True

    def lock_comp(self, node_ids, gpu_ids):
        for nid in node_ids:
            for gid in gpu_ids:
                self._nodes[nid].get_gpus()[gid].lock()

    def lock_comm(self, node_ids):
        for nid in node_ids:
            self._nodes[nid].lock()

    def free_comp(self, node_ids, gpu_ids):
        for nid in node_ids:
            for gid in gpu_ids:
                self._nodes[nid].get_gpus()[gid].free()
        return True

    def free_comm(self, node_ids):
        for nid in node_ids:
            self._nodes[nid].free()

    def get_num_nodes(self):
        return len(self._nodes)

    def find_nodes(self, job):
        requested_gpus = job._ngpus
        match_nodes = []
        for node in self._nodes:
            if node._ngpus >= requested_gpus:
                match_nodes.append(node)
        # If requested_gpus is larger than # of GPUs on a single node
        if len(match_nodes) == 0:
            match_nodes = self._nodes
        return match_nodes


class Node:

    def __init__(self, node_id, node_name, gpus):

        self._node_id = node_id
        self._node_name = node_name
        self._cpu_mem = 64 * (2**30) 
        self._gpus = gpus
        self._ngpus = len(gpus)
        self._net_spd = 10 * (2**30)

        self._net_free = True

        for gpu in self._gpus:
            gpu.set_node(self)

    def __str__(self):
        return '%s:%d,%d' % (self._node_name, self._ngpus, self.workload)

    def __repr__(self):
        return str(self)

    def get_gpus(self):
        return self._gpus

    def get_node_id(self):
        return self._node_id

    def get_status(self):
        return self._net_free

    def lock(self):
        self._net_free = False

    def free(self):
        self._net_free = True

    @property
    def workload(self):
        job_set = set()
        for gpu in self._gpus:
            for j in gpu.get_jobs():
                job_set.add(j._job_id)
        return len(job_set)


class GPU:
    def __init__(self, node_id, node_name, gpu_id, idx):
        
        self._node_id = node_id
        self._node_name = node_name
        self._node = None
        self._gpu_id = gpu_id
        self._idx = idx
        self._gpu_mem = 11 * (2**30) * 0.9

        self._free_mem = self._gpu_mem
        self._alloc_mem = {}

        self._job_list = []
        self._job_map = {}
        self._is_free = True
        self._running_job = None

    def __str__(self):
        return 'node%d-GPU%s:%d/%d MB: %s' % (self._node_id, self._gpu_id, self.get_allocated_mem()/2**20, self._gpu_mem/2**20, self.get_jobs())

    def __repr__(self):
        return str(self)

    def print_mem(self):
        print("GPU %d-%d: %.2f/%.2f MB" % (self._node_id, self._gpu_id, self.get_allocated_mem() / 2**20, self._gpu_mem / 2**20))

    def get_gpu_idx(self):
        return self._idx

    def get_gpu_id(self):
        return self._gpu_id

    def get_node_id(self):
        return self._node.get_node_id()

    def set_node(self, node):
        self._node = node

    def get_node(self):
        return self._node

    def run_job(self, job):
        self._running_job = job
        self._is_free = False

    def end_job(self):
        self._running_job = None
        self._is_free = True

    def get_status(self):
        return self._is_free

    def lock(self):
        self._is_free = False

    def free(self):
        self._is_free = True

    def add_job(self, job):
        
        self._job_list.append(job)
        self._job_map[job._job_id] = job
        self._alloc_mem[job._job_id] = job.get_occupied_mem()

    def remove_job(self, job):
        self._job_list.remove(job)
        self._job_map.pop(job._job_id)
        self._alloc_mem.pop(job._job_id)

    def get_total_mem(self):
        return self._gpu_mem

    def get_free_mem(self):
        alloc_mem = sum(self._alloc_mem.values())
        return self._gpu_mem - alloc_mem

    def get_allocated_mem(self):
        alloc_mem = sum(self._alloc_mem.values())
        return alloc_mem

    def get_workload(self):
        return len(self._job_list)

    def get_jobs(self):
        return self._job_list

    def get_job(self, job_id):
        return self._job_map.get(job_id, None)

    def update_allocated_mem(self, job_id, new_mem):
        self._alloc_mem[job_id] = new_mem
